# -*- coding: utf-8 -*-
import numpy as np
from node import Node
from ElementBase import elementBase
from tools import *
from typing import List
from multipledispatch import dispatch
from Loads import *

class PS3(elementBase):
    """平面应力单元,3节点,线性,节点自由度:Ux,Uy"""
    VTK_TYPE = 5
    elementType = "PlaneStresTri"
    nodeVars=["U1",'U2']
    SupportBCTypes=["DisplacementBc",]
    SupportLoadTypes=['ConcentratedLoad','BodyLoad',"PressureLoad","SurfaceLoad"]
    # 单元形函数
    N1=lambda xiVal,etaVal:1-xiVal-etaVal
    N2=lambda xiVal,etaVal:xiVal
    N3=lambda xiVal,etaVal:etaVal
    # 形函数对xi偏导
    N1xi=-1
    N2xi=1
    N3xi=0
    # 形函数对eta求导
    N1eta=-1
    N2eta=0
    N3eta=1
    def __init__(self, label:int ,
                 node1:Node,
                 node2:Node,
                 node3:Node,
                 density:float=0.0,E:float=0.0,
                 nu:float=0.0,thickness:float=0.0):
        super().__init__(label, 2)
        self.node1=node1
        self.node2=node2
        self.node3=node3
        self.density=density
        self.E=E
        self.nu=nu
        self.thickness=thickness
        # 初始化节点位移/力 列向量
        # [U1x,U1y,U2x,U2y,U3x,U3y]^T
        self._U=np.zeros((6,1))
        self._F=np.zeros((6,1))
    @property
    def get_connectivity(self)->List[int]:
        return [self.node1.label,self.node2.label,self.node3.label]
    @property
    def Ke(self)->np.ndarray:
        """计算单元刚度矩阵,使用单高斯积分点计算"""
        return self.area*self.thickness*(self.B.T.dot(self.D).dot(self.B))
    @property
    def Me(self)->np.ndarray:
        """计算单元质量矩阵,使用3个高斯积分点计算"""
        xi1,xi2,xi3=1.0/6.0,2.0/3.0,1.0/6.0
        eta1,eta2,eta3=1.0/6.0,1.0/6.0,2.0/3.0
        w1,w2,w3=1.0/6.0,1.0/6.0,1.0/6.0
        tmp=w1*self.ShapeFun_at_xi_eta(xi1,eta1).T.dot(self.ShapeFun_at_xi_eta(xi1,eta1))*np.linalg.det(self.J)+\
                w2*self.ShapeFun_at_xi_eta(xi2,eta2).T.dot(self.ShapeFun_at_xi_eta(xi2,eta2))*np.linalg.det(self.J)+\
                w3*self.ShapeFun_at_xi_eta(xi3,eta3).T.dot(self.ShapeFun_at_xi_eta(xi3,eta3))*np.linalg.det(self.J)
        return self.thickness*self.density*tmp
    @property
    def area(self)->float:
        """计算单元面积,使用海伦公式"""
        L12=np.sqrt((self.node2.x-self.node1.x)**2+(self.node2.y-self.node1.y)**2)
        L23=np.sqrt((self.node3.x-self.node2.x)**2+(self.node3.y-self.node2.y)**2)
        L31=np.sqrt((self.node1.x-self.node3.x)**2+(self.node1.y-self.node3.y)**2)
        p=0.5*(L12+L23+L31)
        return np.sqrt(p*(p-L12)*(p-L23)*(p-L31))
    @property
    def D(self)->np.ndarray:
        """材料刚度矩阵"""
        E,nu=self.E,self.nu
        D=np.array([[1,nu,0],
                    [nu,1,0],
                    [0,0,(1 -nu)/2]])*E/(1-nu**2)
        return D
    def ShapeFun_at_xi_eta(self,xi:float,eta:float)->np.ndarray:
        """计算形函数N(2x6 matrix)在(xi,eta)处的值"""
        xiVal,etaVal=xi,eta
        N1=PS4.N1(xiVal,etaVal)
        N2=PS4.N2(xiVal,etaVal)
        N3=PS4.N3(xiVal,etaVal)
        return np.array([[N1,0,N2,0,N3,0],
                         [0,N1,0,N2,0,N3,0]])
    @property
    def J(self)->np.ndarray:
        """计算Jacobian矩阵,对于PS3,Jacobian矩阵为常数矩阵"""
        dx21=self.node2.x-self.node1.x
        dy21=self.node2.y-self.node1.y
        dx31=self.node3.x-self.node1.x
        dy31=self.node3.y-self.node1.y
        return np.array([[dx21,dy21],
                        [dx31,dy31]])
    @property
    def Jinv(self)->np.ndarray:
        """雅可比求逆"""
        return np.linalg.inv(self.J)
    @property
    def B(self)->np.ndarray:
        """计算B矩阵,对于PS3,B矩阵为常数矩阵"""
        N1x,N1y=np.linalg.solve(self.J,np.array([PS3.N1xi,PS3.N1eta]))
        N2x,N2y=np.linalg.solve(self.J,np.array([PS3.N2xi,PS3.N2eta]))
        N3x,N3y=np.linalg.solve(self.J,np.array([PS3.N3xi,PS3.N3eta]))
        return np.array([[N1x,0.0,N2x,0.0,N3x,0.0],
                         [0.0,N1y,0.0,N2y,0.0,N3y],
                         [N1y,N1x,N2y,N2x,N3y,N3x]])
    @property
    def dofNum(self)->int:
        """单元自由度数"""
        return 6
    @property
    def NodesList(self)->List[Node]:
        """获取当前单元的节点列表"""
        return [self.node1,self.node2,self.node3]
    @property
    def ElemDofIndexs(self)->List[int]:
        """获取当前单元的各节点自由度的全局编号"""
        return dofInds(self.node1.label,self.NodeDofNum)+dofInds(self.node2.label,self.NodeDofNum)+dofInds(self.node3.label,self.NodeDofNum)
    # !--> 将load转换为等效节点力
    @dispatch(BodyLoad)
    def ProcessELoad2EqNodeForce(self, load:BodyLoad):
        """将体力转换为等效节点力"""
        eqforce=np.zeros((6,1))
        bodyForceArray=np.zeros((2,1))
        if self.label not in load.get_ElemsList:
            return eqforce
        if load.Direction == 'x':
            bodyForceArray[0,0]=load.LoadValue
        elif load.Direction == 'y':
            bodyForceArray[1,0]=load.LoadValue
        else:
            raise ValueError("BodyLoad direction should be 'x' or 'y'!")
        eqforce=0.5*self.thickness*np.linalg.det(self.J)*self.ShapeFun_at_xi_eta(1.0/3.0,1.0/3.0).T.dot(bodyForceArray)
        return eqforce        
    @dispatch(PressureLoad)
    def processELoad2EqNodeForce(self,load:PressureLoad)->np.ndarray:
        """将压力载荷转为节点力"""
        eqforce=np.zeros((6,1))
        if self.label not in load.get_ElemsList:
            return eqforce
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(load.get_ElemsList, [self.label]))[0]
        for pos in positions:
            if load.FaceIdList[pos] == 1:
                eqforce+=self.SurfLoadOnEdge(load.FaceIdList[pos],self.Face1Normal,load.PressureValue)
            if load.FaceIdList[pos] == 2:
                eqforce+=self.SurfLoadOnEdge(load.FaceIdList[pos],self.Face2Normal,load.PressureValue)
            if load.FaceIdList[pos] == 3:
                eqforce+=self.SurfLoadOnEdge(load.FaceIdList[pos],self.Face3Normal,load.PressureValue)
        return eqforce
    @dispatch(SurfaceLoad)
    def processELoad2EqNodeForce(self,load:SurfaceLoad)->np.ndarray:
        """将面载荷转为节点力"""
        eqforce=np.zeros((6,1))
        if self.label not in load.get_ElemsList:
            return eqforce
        # 找出这个单元标签的位置在load.get_ElemsList中的所有位置
        positions = np.where(np.in1d(load.get_ElemsList, [self.label]))[0]
        for pos in positions:
                eqforce+=self.SurfLoadOnEdge(load.FaceIdList[pos],load.DirectionVector,load.LoadValue)
        return eqforce
    def SurfLoadOnEdge(self,FaceId:int,LoadDirection:np.ndarray,PressureValue:float)->np.ndarray:
        """计算作用在单元边界的外力"""
        if FaceId == 1:
            Tx,Ty,_=LoadDirection*PressureValue
            Len12=np.sqrt((self.node2.x-self.node1.x)**2+(self.node2.y-self.node1.y)**2)
            eqforce=np.array([[Tx],[Ty],[Tx],[Ty],[0.0],[0.0]])*0.5*self.thickness*Len12
        if FaceId == 2:
            Tx,Ty,_=LoadDirection*PressureValue
            Len23=np.sqrt((self.node3.x-self.node2.x)**2+(self.node3.y-self.node2.y)**2)
            eqforce=np.array([[0.0],[0.0],[Tx],[Ty],[Tx],[Ty]])*0.5*self.thickness*Len23
        if FaceId == 3:
            Tx,Ty,_=LoadDirection*PressureValue
            Len13=np.sqrt((self.node1.x-self.node3.x)**2+(self.node1.y-self.node3.y)**2)
            eqforce=np.array([[Tx],[Ty],[0.0],[0.0],[Tx],[Ty]])*0.5*self.thickness*Len13
        return eqforce # type: ignore
    @property
    def Normal(self)->np.ndarray:
        """单元法向, return (3,) array"""
        vec12=np.array([self.node2.x-self.node1.x,self.node2.y-self.node1.y])
        vec13=np.array([self.node3.x-self.node1.x,self.node3.y-self.node1.y])
        return np.cross(vec12,vec13)
    @property
    def Face1Normal(self)->np.ndarray:
        """面1(1-2)外法向,return (3,) array"""
        vec12=np.array([self.node2.x-self.node1.x,self.node2.y-self.node1.y])
        return np.cross(vec12,self.Normal)
    @property
    def Face2Normal(self)->np.ndarray:
        """面2(2-3)外法向,return (3,) array"""
        vec23=np.array([self.node3.x-self.node2.x,self.node3.y-self.node2.y])   
        return np.cross(vec23,self.Normal)
    @property
    def Face3Normal(self)->np.ndarray:
        """面3(3-1)外法向,return (3,) array"""
        vec31=np.array([self.node1.x-self.node3.x,self.node1.y-self.node3.y])
        return np.cross(vec31,self.Normal)
